package com.tensquare.ai.service;

import com.tensquare.ai.util.CnnUtil;
import org.deeplearning4j.models.embeddings.loader.WordVectorSerializer;
import org.deeplearning4j.models.embeddings.wordvectors.WordVectors;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Service;
import util.IKUtil;

import java.io.File;
import java.io.IOException;
import java.util.Random;
import java.util.Map;

/**
 * 构建卷积神经网络模型
 */
@Service
public class CnnService
{

    @Value("${ai.wordLib}")
    private String wordLib; //合并后的分词语料库文件

    @Value("${ai.savePath}")
    private String savePath;//合并前的分词语料库目录

    @Value("${ai.vecModel}")
    private String vecModel;//词向量文件

    @Value("${ai.cnnModel}")
    private String cnnModel;//卷积神经模型文件

    public void build()
    {

        //1.创建计算图对象（封装卷积神经网络算法）
        ComputationGraph cg = CnnUtil.createComputationGraph(100);

        //2.创建数据集对象（封装训练所需数据）
        String[] childPaths = {"ai", "db", "web"};
        //读取回词向量模型
        WordVectors wordVectors = WordVectorSerializer.loadStaticModel(new File(vecModel));
        DataSetIterator dataSetIterator = CnnUtil.getDataSetIterator(savePath, childPaths, wordVectors, 100, 200, new Random(40));

        //3.训练（算法+数据），得到卷积神经网络模型
        cg.fit(dataSetIterator);

        //4.把训练结果保存成文件
        try
        {
            ModelSerializer.writeModel(cg, cnnModel, true);
        } catch (IOException e)
        {
            e.printStackTrace();
        }
    }

    /**
     * 使用卷积神经网络模型进行分类预测
     */
    public Map<String, Double> yuce(String content)
    {
        //先把文章内容进行分词，方便模型进行匹配
        try
        {
            content = IKUtil.split(content, " ");
        } catch (IOException e)
        {
            e.printStackTrace();
        }

        String[] childPaths = {"ai", "db", "web"};
        try
        {
            return CnnUtil.predictions(vecModel, cnnModel, savePath, childPaths, content);
        } catch (IOException e)
        {
            e.printStackTrace();
        }
        return null;
    }

}
